#!/bin/bash

declare -A epochs=(["mrpc"]=30 ["qnli"]=25 ["rte"]=80 ["sst2"]=60 ["stsb"]=40 ["cola"]=80)
declare -A bs=(["mrpc"]=64 ["qnli"]=64 ["rte"]=64 ["sst2"]=64 ["stsb"]=64 ["cola"]=64)
declare -A ml=(["mrpc"]=256 ["qnli"]=256 ["rte"]=512 ["sst2"]=256 ["stsb"]=256 ["cola"]=256)
declare -A lr=(["mrpc"]="4e-4" ["qnli"]="4e-4" ["rte"]="4e-4" ["sst2"]="5e-4" ["stsb"]="4e-4" ["cola"]="4e-4")
declare -A metrics=(["mrpc"]="accuracy" ["qnli"]="accuracy" ["rte"]="accuracy" ["sst2"]="accuracy" ["stsb"]="pearson" ["cola"]="matthews_correlation")
declare -A target=(["all"]="query value key attention.output.dense intermediate.dense output.dense" ["qv"]="query value")
declare -A eval_steps_list=(["Ada"]=100 ["Mask"]=50)

export WANDB_MODE=offline

export CUBLAS_WORKSPACE_CONFIG=":16:8" # https://docs.nvidia.com/cuda/cublas/index.html#cublasApi_reproducibility
export PYTHONHASHSEED=0
export output_dir="./cola"

run(){

  python examples/text-classification/run_glue.py \
  --model_name_or_path roberta-base \
  --task_name cola \
  --do_train \
  --do_eval \
  --max_seq_length 512 \
  --per_device_train_batch_size 32 \
  --learning_rate 4e-4 \
  --num_train_epochs 80 \
  --output_dir $output_dir/model \
  --overwrite_output_dir \
  --logging_steps 10 \
  --logging_dir $output_dir/log \
  --evaluation_strategy epoch \
  --save_strategy epoch \
  --warmup_ratio 0.06 \
  --apply_lora \
  --lora_r 8 \
  --lora_alpha 16 \
  --seed 0 \
  --weight_decay 0.1

}
run "$task" "$rank" "$mode" "$seed" "$Iidea" "$RS_flag"

# ========== 以下代码已被注释掉 ==========
: <<'END_COMMENT'
#task_base=('mrpc' 'qnli' 'rte' 'sst2' 'stsb' 'cola')
task_base=('cola')
seeds=(6 10 11 15 16 20 21 25)
#init_rank=(4 6 12)
#end_rank=(2 4 8)
ranks=(4)
# 1 is Ada, 2 is Mask
modes=("Ada")
# 0 is none, 1 is fixexp_sigmoid, 2 is expnorm
Iideas=("none")
RS=(True)

for task in "${task_base[@]}"; do
  for seed in "${seeds[@]}"; do
    for rank in "${ranks[@]}"; do
      for RS_flag in "${RS[@]}"; do
        for mode in "${modes[@]}"; do
          if [ "$mode" = "Ada" ]; then
            Iidea="none"
            echo "Running with task: $task, seed: $seed, rank: $rank, mode: $mode, Iidea: $Iidea, RS_flag: $RS_flag"
            run "$task" "$rank" "$mode" "$seed" "$Iidea" "$RS_flag"
          else
            for Iidea in "${Iideas[@]}"; do
              echo "Running with task: $task, seed: $seed, rank: $rank, mode: $mode, Iidea: $Iidea, RS_flag: $RS_flag"
              run "$task" "$rank" "$mode" "$seed" "$Iidea" "$RS_flag"
            done
          fi
        done
      done
    done
  done
done
END_COMMENT